//
// Copyright (C) 2004 Andras Varga
//
// This program is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public License
// as published by the Free Software Foundation; either version 2
// of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with this program; if not, see <http://www.gnu.org/licenses/>.
//

#ifndef __INET_IP_H
#define __INET_IP_H

#include "QueueBase.h"
#include "InterfaceTableAccess.h"
#include "RoutingTableAccess.h"
#include "IRoutingTable.h"
#include "ICMPAccess.h"
#include "IPControlInfo.h"
#include "IPDatagram.h"
#include "IPFragBuf.h"
#include "ProtocolMap.h"

#include<queue>


class ARPPacket;
class ICMPMessage;

// ICMP type 2, code 4: fragmentation needed, but don't-fragment bit set
const int ICMP_FRAGMENTATION_ERROR_CODE = 4;


/**
 * Implements the IP protocol.
 */
class INET_API IP : public QueueBase
{
  protected:
    IRoutingTable *rt;
    IInterfaceTable *ift;
    ICMPAccess icmpAccess;
    cGate *queueOutGate; // the most frequently used output gate

    // config
    int defaultTimeToLive;
    int defaultMCTimeToLive;
    simtime_t fragmentTimeoutTime;

    // working vars
    long curFragmentId; // counter, used to assign unique fragmentIds to datagrams
    IPFragBuf fragbuf;  // fragmentation reassembly buffer
    simtime_t lastCheckTime; // when fragbuf was last checked for state fragments
    ProtocolMapping mapping; // where to send packets after decapsulation

    // statistics
    int numMulticast;
    int numLocalDeliver;
    int numDropped;
    int numUnroutable;
    int numForwarded;

    // threshold
    double threshold_light; // threshold for light flow
    double threshold_medium; // threshold for medium flow
    double threshold_reset; // if the link's load falls below this, we reset the table
    double threshold_congestion; // threshold for congestion

    // congestion count
    std::map<IPAddress, int> conSigMap; // for sending out Congestion signal

    // For flow classification
    class FlowKey {
    private:
    	IPAddress src;
    	IPAddress dst;
    	int flowID; // uniquely identify a flow from src -> dst
    	short trafficType;
    	char priority;
    public:
    	FlowKey(IPAddress src_, IPAddress dst_, int flowID_, short trafficType_, char priority_)
			: src(src_), dst(dst_), flowID(flowID_), trafficType(trafficType_), priority(priority_) {
    	}

        IPAddress getDst() const
        {
            return dst;
        }

        int getFlowID() const
        {
            return flowID;
        }

        char getPriority() const
        {
            return priority;
        }

        IPAddress getSrc() const
        {
            return src;
        }

        short getTrafficType() const
        {
            return trafficType;
        }

        void setDst(IPAddress dst)
        {
            this->dst = dst;
        }

        void setFlowID(int flowID)
        {
            this->flowID = flowID;
        }

        void setPriority(char priority)
        {
            this->priority = priority;
        }

        void setSrc(IPAddress src)
        {
            this->src = src;
        }

        void setTraficType(short  trafficType)
        {
            this->trafficType = trafficType;
        }

        bool operator<(const FlowKey &right) const {
        	if (src < right.src)
        		return true;
        	if ((src == right.src) && (dst < right.dst))
        		return true;
        	if ((src == right.src) && (dst == right.dst) && (flowID < right.flowID))
        		return true;
        	if ((src == right.src) && (dst == right.dst) && (flowID == right.flowID) &&
        			(priority < right.priority))
        		return true;
        	if ((src == right.src) && (dst == right.dst) && (flowID == right.flowID) &&
        			(priority == right.priority) && (trafficType < right.trafficType))
        		return true;

        	return false;
		}

        bool operator==(const FlowKey &right) const {
        	return (src == right.src && dst == right.dst &&
        			flowID == right.flowID && priority == right.priority && trafficType == right.trafficType);
        }

        bool operator!=(const FlowKey &right) const {
        	return !(operator==(right));
        }

        FlowKey& operator=(const FlowKey& other) {
        	if (this != &other) { // make sure not the same object
        		src = other.src;
        		dst = other.dst;
        		flowID = other.flowID;
        		priority = other.priority;
        		trafficType = other.trafficType;
        	}
        	return *this;
        }
    };

    // Interface List
    std::vector<InterfaceEntry *> ifList;

    // Flow Entry Map
    typedef std::map<FlowKey, int> FlowMap;
    FlowMap flowMap;

    // Flow information table
    // FlowTable[interface #][dest edge switch][type of traffic][priority][load level]
    bool***** flowTable;

    // Caching the information of nextHop
    std::map<InterfaceEntry *, IPAddress> nextHopCache;

    // Buffering the packets when we wait for the discovery packet
    typedef std::queue<IPDatagram *> FlowQueue;
    std::map<FlowKey, FlowQueue*> flowQueueCache;
    std::map<FlowKey, IPDatagram *> disPacketCache;

  protected:
    // utility: look up interface from getArrivalGate()
    virtual InterfaceEntry *getSourceInterfaceFrom(cPacket *msg);

    // utility: show current statistics above the icon
    virtual void updateDisplayString();

    /**
     * Encapsulate packet coming from higher layers into IPDatagram, using
     * the control info attached to the packet.
     */
    virtual IPDatagram *encapsulate(cPacket *transportPacket, InterfaceEntry *&destIE);

    /**
     * Encapsulate packet coming from higher layers into IPDatagram, using
     * the given control info. Override if you subclassed controlInfo and/or
     * want to add options etc to the datagram.
     */
    virtual IPDatagram *encapsulate(cPacket *transportPacket, InterfaceEntry *&destIE, IPControlInfo *controlInfo);

    /**
     * Creates a blank IP datagram. Override when subclassing IPDatagram is needed
     */
    virtual IPDatagram *createIPDatagram(const char *name);

    /**
     * Handle IPDatagram messages arriving from lower layer.
     * Decrements TTL, then invokes routePacket().
     */
    virtual void handlePacketFromNetwork(IPDatagram *datagram);

    /**
     * Handle messages (typically packets to be send in IP) from transport or ICMP.
     * Invokes encapsulate(), then routePacket().
     */
    virtual void handleMessageFromHL(cPacket *msg);

    /**
     * Handle incoming ARP packets by sending them over "queueOut" to ARP.
     */
    virtual void handleARP(ARPPacket *msg);

    /**
     * Handle incoming ICMP messages.
     */
    virtual void handleReceivedICMP(ICMPMessage *msg);

    /**
     * Performs routing. Based on the routing decision, it dispatches to
     * reassembleAndDeliver() for local packets, to fragmentAndSend() for forwarded packets,
     * to handleMulticastPacket() for multicast packets, or drops the packet if
     * it's unroutable or forwarding is off.
     */
    virtual void routePacket(IPDatagram *datagram, InterfaceEntry *destIE, bool fromHL);

    /**
     * Performs hash routing.
     */
    virtual void hashRoutePacket(IPDatagram *datagram, bool fromHL);

    /**
     * Forwards packets to all multicast destinations, using fragmentAndSend().
     */
    virtual void routeMulticastPacket(IPDatagram *datagram, InterfaceEntry *destIE, InterfaceEntry *fromIE);

    /**
     * Perform reassembly of fragmented datagrams, then send them up to the
     * higher layers using sendToHL().
     */
    virtual void reassembleAndDeliver(IPDatagram *datagram);

    /**
     * Decapsulate and return encapsulated packet after attaching IPControlInfo.
     */
    virtual cPacket *decapsulateIP(IPDatagram *datagram);

    /**
     * Fragment packet if needed, then send it to the selected interface using
     * sendDatagramToOutput().
     */
    virtual void fragmentAndSend(IPDatagram *datagram, InterfaceEntry *ie, IPAddress nextHopAddr);

    /**
     * Last TTL check, then send datagram on the given interface.
     */
    virtual void sendDatagramToOutput(IPDatagram *datagram, InterfaceEntry *ie, IPAddress nextHopAddr);

  public:
    IP() {}
    ~IP();

  protected:
    /**
     * Initialization
     */
    virtual void initialize();

    /**
     * Processing of IP datagrams. Called when a datagram reaches the front
     * of the queue.
     */
    virtual void endService(cPacket *msg);

    /**
     * Allocates data for the flow table
     */
    virtual void allocateFlowTable();

    /**
     * Update the entries in the flow table
     */
    virtual void updateFlowTableEntry(InterfaceEntry *ie);

    /**
     * get the blocking bit for certain flow
     */
    virtual bool isBlocked(const InterfaceEntry *ie, const IPDatagram *datagram);

    virtual FlowKey makeKey(const IPDatagram *datagram);
    virtual void addFlowEntry(const FlowKey& key, const InterfaceEntry *ie);

    /**
     * Routes a packet in an existing flow
     */
    virtual void hrHandleExistingFlow(IPDatagram *datagram, const int idx);

    /**
     * Routes a packet in a new flow at an edge router
     */
    virtual void hrEdgeNewFlow(IPDatagram *datagram);

    /**
     * The discovery packet (also the first packet) of a flow
     * arrives
     */
    virtual void hrHandleDisPacket(IPDatagram *datagram);

    /**
     * Routes a packet in a new flow at a normal router
     */
    virtual void hrHandleNoDiscovery(IPDatagram *datagram);

    /**
	 * Re-route a flow
	 */
	virtual void hrReroute(IPDatagram *datagram);

    /**
     * Finds a hash route
     **/
    virtual const IPRoute *findHashRoute(const IPDatagram *data, const std::set<IPAddress> *prevHops = NULL);

    /**
     * Handle NACK packet
     */
    virtual void handleNACK(const IPDatagram *datagram);

    /**
     * Handle congestion signal packet
     */
    virtual void handleCongestionSignal(IPDatagram *datagram);

    /**
     * Locate the index of the interface entry to send out a packet
     */
    virtual int getInterfaceIndex(const FlowKey& key);

    /**
     * Sends feedback signal when the router experiences heavy loads
     */
    virtual void sendLoadSignal(const IPDatagram *datagram);
};

#endif

